import logging
from py2neo import Graph
import networkx as nx
import warnings

# 禁用可能的警告信息
warnings.filterwarnings("ignore")

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("PageRankCalculator")


class PageRankCalculator:
    def __init__(self, neo4j_uri: str, neo4j_user: str, neo4j_password: str):
        try:
            self.graph = Graph(neo4j_uri, auth=(neo4j_user, neo4j_password))
            self.graph.run("RETURN 1")
            logger.info("Neo4j 连接成功")
        except Exception as e:
            logger.critical(f"Neo4j连接失败: {e}")
            raise

    def fetch_graph_data(self):
        """加载所有核心知识点节点：仅匹配标签为 KnowledgePoint 或 SubKnowledgePoint 的节点"""
        node_query = """
        MATCH (n)
        WHERE n:KnowledgePoint OR n:SubKnowledgePoint
        RETURN 
            CASE 
                WHEN n:KnowledgePoint THEN n.kp_id
                WHEN n:SubKnowledgePoint THEN n.sub_kp_id
            END AS business_id,
            n
        """
        nodes = {}
        for record in self.graph.run(node_query):
            business_id = record["business_id"]
            nodes[business_id] = record["n"]

        # 边查询：只匹配两端都是 KnowledgePoint 或 SubKnowledgePoint 的关系
        edge_query = """
        MATCH (a)-[r]->(b)
        WHERE (a:KnowledgePoint OR a:SubKnowledgePoint)
          AND (b:KnowledgePoint OR b:SubKnowledgePoint)
        RETURN 
            CASE 
                WHEN a:KnowledgePoint THEN a.kp_id
                WHEN a:SubKnowledgePoint THEN a.sub_kp_id
            END AS source,
            CASE 
                WHEN b:KnowledgePoint THEN b.kp_id
                WHEN b:SubKnowledgePoint THEN b.sub_kp_id
            END AS target
        ORDER BY source, target 
        """
        edges = []
        for record in self.graph.run(edge_query):
            edges.append((record["source"], record["target"]))

        logger.info(f"稳定加载节点数: {len(nodes)}, 边数: {len(edges)}")
        return nodes, edges

    def compute_pagerank(self, top_n: int = 10, damping_factor: float = 0.85):
        """
        构建 NetworkX 图并计算 PageRank，
        同时将结果写入 Neo4j（使用业务ID匹配）。
        """
        nodes, edges = self.fetch_graph_data()
        G = nx.DiGraph()
        G.add_nodes_from(nodes.keys())
        G.add_edges_from(edges)

        initial_pr = {node: 1.0 / len(nodes) for node in G.nodes()}
        pr = nx.pagerank(G, alpha=damping_factor, max_iter=5000, tol=1e-15, nstart=initial_pr)
        self.write_pagerank_to_neo4j(pr)

        sorted_pr = sorted(pr.items(), key=lambda x: x[1], reverse=True)
        result = []
        for node_id, score in sorted_pr[:top_n]:
            node = nodes[node_id]
            name = node.get("title") or node.get("name") or node.get("kp_id") or node.get("sub_kp_id") or str(node_id)
            result.append((name, score))
        return result

    def write_pagerank_to_neo4j(self, pagerank_dict):
        """
        将计算得到的 PageRank 值写入 Neo4j，
        使用 kp_id 或 sub_kp_id 作为业务ID匹配节点。
        """
        try:
            query = """
            UNWIND $pagerankData AS row
            MATCH (n)
            WHERE ( (n:KnowledgePoint AND n.kp_id = row.id)
                    OR (n:SubKnowledgePoint AND n.sub_kp_id = row.id) )
            SET n.pagerank = row.score
            """
            data = [{"id": node_id, "score": score} for node_id, score in pagerank_dict.items()]
            self.graph.run(query, pagerankData=data)
            logger.info("PageRank 值已成功写入 Neo4j 数据库中")
        except Exception as e:
            logger.error(f"写入 PageRank 值失败: {e}")
            raise


if __name__ == "__main__":
    neo4j_uri = "bolt://localhost:7687"
    neo4j_user = "neo4j"
    neo4j_password = "123456789"

    pr_calc = PageRankCalculator(neo4j_uri, neo4j_user, neo4j_password)
    top_results = pr_calc.compute_pagerank(top_n=10)

    print("PageRank 计算结果（前10）：")
    for name, score in top_results:
        print(f"节点: {name}, PageRank: {score}")
